.globl matmul

.text
# =======================================================
# FUNCTION: Matrix Multiplication of 2 integer matrices
# 	d = matmul(m0, m1)
#   The order of error codes (checked from top to bottom):
#   If the dimensions of m0 do not make sense,
#   this function exits with exit code 2.
#   If the dimensions of m1 do not make sense,
#   this function exits with exit code 3.
#   If the dimensions don't match,
#   this function exits with exit code 4.
# Arguments:
# 	a0 (int*)  is the pointer to the start of m0
#	a1 (int)   is the # of rows (height) of m0
#	a2 (int)   is the # of columns (width) of m0
#	a3 (int*)  is the pointer to the start of m1
# 	a4 (int)   is the # of rows (height) of m1
#	a5 (int)   is the # of columns (width) of m1
#	a6 (int*)  is the pointer to the the start of d
# Returns:
#	None (void), sets d = matmul(m0, m1)
# =======================================================
matmul:

    # Error checks
    slti t0, a1, 1
    slti t1, a2, 1
    or t0, t0, t1
    # if t0 == 0, goto check_m1
    beq t0, x0, check_m1
    # call exit2, status code = 2
    addi a1, x0, 2
    jal ra, exit2

check_m1:
    slti t0, a4, 1
    slti t1, a5, 1
    or t0, t0, t1
    # if t0 == 0, goto check_mat
    beq t0, x0, check_mat
    # call exit2, status code = 3
    addi a1, x0, 3
    jal ra, exit2

check_mat:
    beq a2, a4, prologue
    # call exit2, status code = 4
    addi a1, x0, 4
    jal ra, exit2

prologue:
    # Prologue
    addi sp, sp, -16
    sw ra, 0(sp)
    sw s0, 4(sp)
    sw s1, 8(sp)
    sw s2, 12(sp)


    add s0, x0, x0
    add s2, x0, x0

outer_loop_start:
    # while(s0 < a1)
    beq s0, a1, outer_loop_end
    # s1 = 0
    add s1, x0, x0


inner_loop_start:
    # while(s1 < a5)
    beq s1, a5, inner_loop_end
    # prologue
    addi sp, sp, -16
    sw a0, 0(sp)
    sw a1, 4(sp)
    sw a3, 8(sp)
    sw a4, 12(sp)
    # a6[s2] = a0[s0] * a3[s1]
    # a0 = a0 + s0*4*a2
    slli t0, s0, 2
    mul t0, t0, a2
    add a0, a0, t0
    # a1 = a3 + s1*4
    slli t1, s1, 2
    add a1, a3, t1
    # a2 = a2(remain the same)
    # a3 = 1
    addi a3, x0, 1
    # a4 = a5
    add a4, a5, x0
    # call dot
    jal ra, dot
    # a6[s2] = a0
    slli t2, s2, 2
    add t2, a6, t2
    sw a0, 0(t2)
    # epilogue
    lw a0, 0(sp)
    lw a1, 4(sp)
    lw a3, 8(sp)
    lw a4, 12(sp)
    addi sp, sp, 16
    # s1++, s2++
    addi s1, s1, 1
    addi s2, s2, 1
    j inner_loop_start


inner_loop_end:
    # s0++
    addi s0, s0, 1
    j outer_loop_start


outer_loop_end:
    # Epilogue
    lw ra, 0(sp)
    lw s0, 4(sp)
    lw s1, 8(sp)
    lw s2, 12(sp)
    addi sp, sp, 16

    ret
